package dfxnftables

import (
	"context"
	"fmt"
	"net"

	"github.com/google/nftables"
	"github.com/zeromicro/go-zero/core/logc"
)

// AddElement 向 nftables 集合中添加 IP 地址元素。
// 参数:
//   - ctx: 上下文，用于控制请求的生命周期。
//   - c: 与 nftables 内核接口通信的连接对象。
//
// 返回值:
//   - error: 操作过程中发生的错误信息，如果操作成功则返回 nil。
func (n *NfSet) AddElement(ctx context.Context, c *nftables.Conn) error {
	// 校验元数据是否有效
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}
	// 获取表族（IPv4 或 IPv6）
	fam := n.Family()
	// 解析对应的 nftables 表
	tbl, err := n.ResolveTable(c)
	if err != nil {
		return err
	}
	// 获取指定名称的集合
	s, err := c.GetSetByName(tbl, n.SetName)
	if err != nil {
		return err
	}
	// 确保集合存在
	if s == nil {
		return fmt.Errorf("set %s not found", n.SetName)
	}
	// 构造要插入的集合元素列表
	var els []nftables.SetElement
	for _, ipstr := range n.SetIPList {
		ip := net.ParseIP(ipstr)
		if ip == nil {
			continue
		}
		// 处理 IPv6 地址
		if fam == nftables.TableFamilyIPv6 {
			ip = ip.To16()
			if ip == nil {
				continue
			}
			els = append(els, nftables.SetElement{Key: []byte(ip)})
		} else {
			// 处理 IPv4 地址
			ip4 := ip.To4()
			if ip4 == nil {
				continue
			}
			els = append(els, nftables.SetElement{Key: []byte(ip4)})
			logc.Debugf(ctx, "正在插入IP到数组:%s", ip4)
		}
	}
	// 如果没有有效的 IP 元素需要处理，则直接返回
	if len(els) == 0 {
		return nil
	}
	// 将构造好的元素批量加入到集合中
	if err := c.SetAddElements(s, els); err != nil {
		return err
	}
	logc.Debugf(ctx, "正在插入IP到nftset:%s", n.SetIPList)
	// 提交更改至内核
	return c.Flush()
}

// DelElement 从 nftables 集合中删除指定的 IP 地址元素。
// 参数:
//   - ctx: 上下文，用于控制请求的生命周期。
//   - c: 与 nftables 内核接口通信的连接对象。
//
// 返回值:
//   - error: 操作过程中发生的错误信息，如果操作成功则返回 nil。
func (n *NfSet) DelElement(ctx context.Context, c *nftables.Conn) error {
	// 校验元数据是否有效
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}
	// 获取表族（IPv4 或 IPv6）
	fam := n.Family()
	// 解析对应的 nftables 表
	tbl, err := n.ResolveTable(c)
	if err != nil {
		return err
	}
	// 获取指定名称的集合
	s, err := c.GetSetByName(tbl, n.SetName)
	if err != nil {
		return err
	}
	// 确保集合存在
	if s == nil {
		return fmt.Errorf("set %s not found", n.SetName)
	}
	// 构造要删除的集合元素列表
	var els []nftables.SetElement
	for _, ipstr := range n.SetIPList {
		ip := net.ParseIP(ipstr)
		if ip == nil {
			continue
		}
		// 处理 IPv6 地址
		if fam == nftables.TableFamilyIPv6 {
			ip = ip.To16()
			if ip == nil {
				continue
			}
			els = append(els, nftables.SetElement{Key: []byte(ip)})
		} else {
			// 处理 IPv4 地址
			ip4 := ip.To4()
			if ip4 == nil {
				continue
			}
			els = append(els, nftables.SetElement{Key: []byte(ip4)})
		}
	}
	// 如果没有有效的 IP 元素需要处理，则直接返回
	if len(els) == 0 {
		return nil
	}
	// 批量从集合中删除元素
	if err := c.SetDeleteElements(s, els); err != nil {
		return err
	}
	// 提交更改至内核
	return c.Flush()
}
